# Date: 2021/10/28
# Usage: 针对报告单图像的分类识别
import cv2
import numpy as np
import tensorflow as tf
import skimage
import ocr.params as params
from ocr.clf_nets import loader
from ocr.clf_nets.googlenet import GoogLeNet_cifar

label_dict = ['cover', 'conclusion', 'project', 'other']


def rescale_im(im):
    """ Pre-process for images
        images are rescaled so that the shorter side = 224
    """
    im = np.array(im)
    h, w = im.shape[0], im.shape[1]
    if h >= w:
        new_w = 224
        im = skimage.transform.resize(im, (int(h * new_w / w), 224),
                                      preserve_range=True)
    else:
        new_h = 224
        im = skimage.transform.resize(im, (224, int(w * new_h / h)),
                                      preserve_range=True)
    return im.astype('uint8')

def clf_predict(image_array):

    ''' ckpt格式模型导入预测 '''
    # image_array = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)
    # image_array = rescale_im(image_array)
    # image_array = np.array(image_array).astype(np.float32)
    # mean_list = [167.7599334716797, 166.42912928263345, 164.6688690185547]
    # for c_id in range(0, image_array.shape[-1]):
    #     image_array[:, :, c_id] = image_array[:, :, c_id] - mean_list[c_id]
    # image_array = image_array[np.newaxis, :]
    #
    # # Read Cifar label into a dictionary
    # label_dict = loader.load_label_dict(dataset='tijian')
    #
    # # Create a testing GoogLeNet model
    # test_model = GoogLeNet_cifar(
    #     n_channel=3, n_class=4, bn=True, sub_imagenet_mean=False)
    # test_model.create_test_model()
    #
    # with tf.Session() as sess:
    #     saver = tf.train.Saver()
    #     sess.run(tf.global_variables_initializer())
    #     saver.restore(sess, '{}inception-cifar-epoch-{}'.format(params.clf_save_path, params.load))
    #     # get prediction results
    #     pred = sess.run(test_model.layers['top_4'],
    #                     feed_dict={test_model.image: image_array})
    #     # print(pred)
    #     # display results
    #     for re_prob, re_label in zip(pred[0], pred[1]):
    #         print('===============================')
    #         for i in range(3):
    #             print('{}: probability: {:.02f}, label: {}'
    #                   .format(i + 1, re_prob[i], label_dict[re_label[i]]))

    ''' pb格式模型导入预测 '''
    with tf.gfile.GFile(params.clf_save_path, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name='')
        isess = tf.InteractiveSession()

    images_placeholder = graph.get_tensor_by_name("image:0")
    feat = graph.get_tensor_by_name("fc_layers/Squeeze:0")

    # import skimage
    # import imageio
    # image = imageio.imread("F://laibo/Data/tijianbaogao/cls/raw/test_batch/3_other_2.jpg", pilmode='RGB')
    # image = np.array(image)
    # h, w = image_array.shape[0], image_array.shape[1]
    # if h >= w:
    #     new_w = 224
    #     image = skimage.transform.resize(image_array, (int(h * new_w / w), 224),
    #                                   preserve_range=True)
    # else:
    #     new_h = 224
    #     image = skimage.transform.resize(image_array, (224, int(w * new_h / h)),
    #                                   preserve_range=True)
    image_array = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)
    image = rescale_im(image_array)
    image = np.array(image).astype(np.float32)
    mean_list = [167.7599334716797, 166.42912928263345, 164.6688690185547]
    for c_id in range(0, image.shape[-1]):
        image[:, :, c_id] = image[:, :, c_id] - mean_list[c_id]
    image = image[np.newaxis, :]
    pred = isess.run([feat], feed_dict={images_placeholder: image})
    print(label_dict[np.argmax(pred[0][0])])

    return label_dict[np.argmax(pred[0][0])]
